import time
import numpy as np
import pandas as pd
from cvzone.FaceMeshModule import FaceMeshDetector
import cv2
from PIL import Image
from torchvision import transforms
import torch
import pyautogui as pag
from torch.utils.tensorboard import SummaryWriter
import keyboard
import warnings
warnings.filterwarnings("ignore")

img_name_index =174
bool_add = True
writer = SummaryWriter("logs")
leftEyeUpPoint = []
leftEyeDownPoint = []
leftEyeLeftPoint = []
leftEyeRightPoint = []
rightEyeUpPoint = []
rightEyeDownPoint = []
rightEyeLeftPoint = []
rightEyeRightPoint = []
# 初始化网络摄像头
# '2' 表示与计算机连接的第三台相机，'0' 通常指内置网络摄像头
cap = cv2.VideoCapture(0)

# 初始化 FaceMeshDetector 对象
# staticMode: 若为真，则仅检测一次，否则每帧都检测
# maxFaces: 最大可检测人脸数量
# minDetectionCon: 检测置信度阈值
# minTrackCon: 跟踪置信度阈值
detector = FaceMeshDetector(staticMode=False, maxFaces=1, minDetectionCon=0.5, minTrackCon=0.5)

# 人眼关键点所在的索引
idList = [22, 23, 24, 26, 110, 130, 157, 158, 159, 160, 161, 243]  #左眼

# 开始循环以持续获取摄像头帧
while True:
    # 从网络摄像头读取当前帧
    # success: 帧是否成功捕获的布尔值
    # img: 当前帧
    success, img = cap.read()

    # 在图像中找到面部网格
    # img: 如设置 draw=True 则更新了面部网格的图像
    # faces: 检测到的面部信息
    img, faces = detector.findFaceMesh(img, draw=False)
    # 如果检测到面部
    if faces:
        # 遍历每个检测到的面部
        for face in faces:
            # font = cv2.FONT_HERSHEY_SIMPLEX
            # for (index, value) in enumerate(face):
            #     cv2.putText(img, str(index), value, font, 0.2, (255, 230, 0), 1, cv2.LINE_AA)
            # 获取眼睛特定点
            leftEyeUpPoint = face[159]
            leftEyeDownPoint = face[23]
            leftEyeLeftPoint = face[130]
            leftEyeRightPoint = face[243]
            rightEyeUpPoint = face[386]
            rightEyeDownPoint = face[253]
            rightEyeLeftPoint = face[463]
            rightEyeRightPoint = face[466]

            # 绘制眼部周围的特征点
            demo1 = [leftEyeUpPoint, leftEyeDownPoint, leftEyeLeftPoint, leftEyeRightPoint]
            demo2 = [rightEyeUpPoint, rightEyeDownPoint, rightEyeLeftPoint, rightEyeRightPoint]
            # for i in demo1:
            #     cv2.circle(img, i, 1, (0, 0, 0), -1)
            # for i in demo2:
            #     cv2.circle(img, i, 1, (0, 0, 0), -1)

            # 框选出眼睛的矩形图
            aa1x = leftEyeLeftPoint[0] - 10
            aa1y = min(leftEyeUpPoint[1], rightEyeUpPoint[1]) - 10
            aa2x = rightEyeRightPoint[0] + 10
            aa2y = max(rightEyeDownPoint[1], leftEyeDownPoint[1]) + 10
            aa1 = [aa1x, aa1y]
            aa2 = [aa2x, aa2y]


            # cv2.rectangle(img,aa1, aa2, (255, 0, 0), 2)

            def cropped_eyes():
                x = aa1x
                y = aa1y
                w = abs(aa1x - aa2x)
                h = abs(aa1y - aa2y)
                return x, y, w, h


            x, y, w, h = cropped_eyes()
            cropped_screen = img[y:y + h, x:x + w]


            def get_image():
                img_begin = torch.from_numpy(cropped_screen)
                image = torch.tensor(img_begin).permute(2, 0, 1)
                image = image.unsqueeze(0)
                return image.float() / 127. - 1.



            def add_data():
                global img_name_index
                global bool_add
                if bool_add:
                    label_data = pd.read_csv('./data/val/label/labels.csv')

                    screenWidth, screenHeight = pag.size()
                    x, y = pag.position()
                    # 返回鼠标的坐标
                    print('屏幕： (%s,%s)  鼠标坐标 : (%s, %s)' % (screenWidth, screenHeight ,x, y))


                    # 使用concat添加新行
                    new_row = pd.DataFrame([[img_name_index,(x/screenWidth)*2-1, (y/screenHeight)*2-1]], columns=['ID','x', 'y'])  # 创建一个只包含新行的DataFrame
                    label_data = pd.concat([label_data, new_row], ignore_index=True)  # 将新行添加到原始DataFrame

                    label_data.to_csv('./data/val/label/labels.csv', index=False)
                    cv2.imwrite("./data/val/img/" + str(img_name_index) + ".png", cropped_screen)

                    img_name_index += 1
                    tensor_img = get_image()
                    print(tensor_img.shape)
                    writer.add_images("eyes", tensor_img)
                    bool_add = False


            # 按下空格收集数据集
            keyboard.add_hotkey(" ", add_data, suppress=False)
            bool_add = True

    # 在名为 'Image' 的窗口中显示图像
    cv2.imshow("Image", img)

    # 等待 1 毫秒检查用户输入，保持窗口打开状态
    cv2.waitKey(1)
